package io.github.wanggit.tracing.agent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.List;

public class TracingClassFileTransformer implements ClassFileTransformer {

    public static final String CGLIB_CLASS_SEPARATOR = "$$";

    public static final String HIBERNATE_PROXY = "$HibernateProxy$";

    private final List<String> pkgPrefixs;

    private final boolean showLog;

    public TracingClassFileTransformer(TracingAgent.Args args) {
        pkgPrefixs = args.getPkcPrefixs();
        showLog = args.isShowLog();
    }

    private boolean needTracing(String className) {
        for (String pkgPrefix : pkgPrefixs) {
            if (className.startsWith(pkgPrefix)) {
                return true;
            }
        }
        return false;
    }

    private boolean needInsertCodes(CtMethod ctMethod, CtClass ctClass) {
        return !ctMethod.isEmpty()
                && !ctMethod.getLongName().startsWith("java.lang.Object.")
                && null != ctMethod.getMethodInfo2().getCodeAttribute()
                && !ctClass.isFrozen()
                && !ctClass.getName().contains(CGLIB_CLASS_SEPARATOR)
                && !ctClass.getName().contains(HIBERNATE_PROXY);
    }

    private void printLog(String message) {
        if (showLog) {
            System.out.println(message);
        }
    }

    @Override
    public byte[] transform(
            ClassLoader loader,
            String className,
            Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain,
            byte[] classfileBuffer)
            throws IllegalClassFormatException {
        if (!needTracing(className)) {
            return classfileBuffer;
        }
        try {
            CtClass ctClass =
                    ClassPool.getDefault().makeClass(new ByteArrayInputStream(classfileBuffer));
            if (!ctClass.isAnnotation()
                    && !ctClass.isInterface()
                    && !ctClass.isEnum()
                    && !ctClass.isPrimitive()
                    && !ctClass.isFrozen()) {
                CtMethod[] methods = ctClass.getDeclaredMethods();
                for (CtMethod ctMethod : methods) {
                    if (needInsertCodes(ctMethod, ctClass)) {
                        printLog(ctMethod.getLongName() + " to transformer...");
                        ctMethod.insertBefore(
                                "{ "
                                        + "io.github.wanggit.antrpc.client.monitor.report.ZipkinAllMethodTracing.start(\""
                                        + className
                                        + "\"); "
                                        + "}");
                        ctMethod.insertAfter(
                                "{ io.github.wanggit.antrpc.client.monitor.report.ZipkinAllMethodTracing.trace(\""
                                        + className
                                        + "\", \""
                                        + ctMethod.getLongName()
                                        + "\"); }");
                    }
                }
                return ctClass.toBytecode();
            }
        } catch (Exception e) {
            printLog("AntRPC Tracing Agent ClassFile Transformer ERROR. ");
        }
        return classfileBuffer;
    }
}
